import os
import openai
from langchain_community.llms import OpenAI
from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate

os.environ["OPENAI_API_KEY"]=os.environ["OPENAI_API_KEY_ZHIHU"]
os.environ["OPENAI_API_BASE"]=os.environ["OPENAI_API_BASE_ZHIHU"]
llm = OpenAI(model_name= "gpt-3.5-turbo",temperature=0.9)

# #示例一
# prompt = """The following is a conversation with an AI assistant.
# The assistant is typically sarcastic and witty, producing creative and funny responses to the users questions.
#
# User: What defines a good network engineer?
# Please answer the question"""
#
# print(llm(prompt))


# #示例二
# prompt = """The following are exerpts from conversations with an AI assistant. The assistant is typically sarcastic and
# witty, producing creative and funny responses to the users questions. Here are some examples:
#
# User: My router is not working
# AI: Would you please reboot the router? Let me give you the cli command.
#
# User: It is still not working
# AI: Would you please run a diagnostic against the line card
#
# User: Can we do it tomorrow?
# AI: No, this is a very urgent issue, please bear with me for a moment
#
# User: What defines a network engineer
# AI:"""
#
# print(llm(prompt))

#示例三
examples = [
    {
        "query":"My router is not working",
        "answer":"Would you please reboot the router? Let me give you the cli command."
    },
    {
        "query":"It is still not working",
        "answer":"Would you please run a diagnostic against the line card"
    },
    {
        "query":"Can we do it tomorrow?",
        "answer":"No, this is a very urgent issue, please bear with me for a moment"
    }
]

examples_template = """
User: {query}
AI: {answer}
"""

examples_prompt = PromptTemplate(input_variables=["query","answer"],template=examples_template)

prefix = """The following are exerpts from conversations with an AI assistant. The assistant is typically sarcastic and
 witty, producing creative and funny responses to the users questions. Here are some examples:"""
suffix = """User: {query}
AI:"""

few_shot_prompt_template = FewShotPromptTemplate(input_variables=["query"],examples = examples,example_prompt = examples_prompt,
                                                 prefix=prefix,suffix=suffix,example_separator="\n\n")

llm_chain = LLMChain(prompt=few_shot_prompt_template,llm=llm)

query = "What defines a network engineer"
print(llm_chain(few_shot_prompt_template.format(query=query))["text"])